import numpy as np
import struct
import os
import faiss


# to get the .vecs
# np.set_printoptions(threshold=np.inf)  # display all the content when print the numpy array


def ivecs_read(fname):
    a = np.fromfile(fname, dtype='int32')
    d = a[0]
    return a.reshape(-1, d + 1)[:, 1:].copy(), d


def fvecs_read(fname):
    data, d = ivecs_read(fname)
    return data.view('float32').astype(np.float32), d


def bvecs_read(fname):
    a = np.fromfile(fname, dtype='uint8')
    d = a[:4].view('uint8')[0]
    return a.reshape(-1, d + 4)[:, 4:].copy(), d


# store in format of vecs
def fvecs_write(filename, vecs):
    f = open(filename, "wb")
    dimension = [len(vecs[0])]

    for x in vecs:
        f.write(struct.pack('i' * len(dimension), *dimension))  # *dimension就是int, dimension就是list
        f.write(struct.pack('f' * len(x), *x))

    f.close()


def ivecs_write(filename, vecs):
    f = open(filename, "wb")
    dimension = [len(vecs[0])]

    for x in vecs:
        f.write(struct.pack('i' * len(dimension), *dimension))
        f.write(struct.pack('i' * len(x), *x))

    f.close()


def bvecs_write(filename, vecs):
    f = open(filename, "wb")
    dimension = [len(vecs[0])]

    for x in vecs:
        f.write(struct.pack('i' * len(dimension), *dimension))
        f.write(struct.pack('B' * len(x), *x))

    f.close()


def bin_read(filename, dimension):
    a = np.fromfile(filename, dtype=np.float32).reshape(-1, dimension)
    return a


def bin_write(filename, vecs):
    f = open(filename, "wb")

    for x in vecs:
        f.write(struct.pack('f' * len(x), *x))

    f.close()


def ip_gnd(base, query, k):
    base_dim = base.shape[1]
    index = faiss.IndexFlatIP(base_dim)
    index.add(base)
    gnd_distance, gnd_idx = index.search(query, k)
    return gnd_idx, gnd_distance


def delete_dir_if_exist(dire):
    if os.path.isdir(dire):
        command = 'rm -rf %s' % dire
        print(command)
        os.system(command)


def fvecs2bin(vecs_basic_dir, bin_basic_dir, dataset_name, topk):
    vecs_base_fname = '%s/%s/%s_base.fvecs' % (vecs_basic_dir, dataset_name, dataset_name)
    vecs_query_fname = '%s/%s/%s_query.fvecs' % (vecs_basic_dir, dataset_name, dataset_name)
    base, dim = fvecs_read(vecs_base_fname)
    query, dim = fvecs_read(vecs_query_fname)

    delete_dir_if_exist("%s/%s" % (bin_basic_dir, dataset_name))
    os.mkdir("%s/%s" % (bin_basic_dir, dataset_name))
    bin_base_fname = '%s/%s/%s_base.bin' % (bin_basic_dir, dataset_name, dataset_name)
    bin_query_fname = '%s/%s/%s_query.bin' % (bin_basic_dir, dataset_name, dataset_name)
    bin_write(bin_base_fname, base)
    bin_write(bin_query_fname, query)
    print("write %s base query" % dataset_name)

    gnd, _ = ip_gnd(base, query, topk)
    bin_gnd_fname = '%s/%s/gnd.ivecs' % (bin_basic_dir, dataset_name)
    ivecs_write(bin_gnd_fname, gnd)
    print("write %s groundtruth" % dataset_name)


if __name__ == '__main__':
    dataset_l = ['audio', 'imagenet', 'movielens', 'music100', 'netflix', 'tiny5m', 'word2vec', 'yahoomusic']
    # dataset_l = ['audio']
    for ds_name in dataset_l:
        fvecs2bin('/home/zhengbian/Dataset/inner_product', '/home/zhengbian/ip-nsw/data', ds_name, 10)
